#!/usr/bin/env Rscript
##################################
############ Preamble ############
##################################

# set language to English
Sys.setenv(LANG = "en")

# clean up
rm(list = ls())

# Set working directory: please set your own
if (Sys.getenv("RSTUDIO") == "1") setwd("~/Dropbox/cues_bjps_replication/")

# Load necessary packages
library(DeclareDesign)
library(DesignLibrary)
library(tidyverse)
library(future.apply)


# N in each group:
# FULL SAMPLE
# Control: 413
# Mainstream and RRP Approve: 414
# Mainstream Approve: 406
# Mainstream Disapprove and RRP Approve: 386
# RRP Approve: 394

# LEFT-WING SAMPLE
# Control: 209
# Mainstream and RRP Approve: 208
# Mainstream Approve: 205
# Mainstream Disapprove and RRP Approve: 200
# RRP Approve: 199

# RIGHT-WING SAMPLE
# Control: 151
# Mainstream and RRP Approve: 138
# Mainstream Approve: 131
# Mainstream Disapprove and RRP Approve: 129
# RRP Approve: 131


# Set seed
set.seed(2004)

# Note: I always use the median N in each subsample
# For right-wing sample
n_full_sample <- c(413 + 414, 413 + 406, 413 + 386, 413 + 394)
n_right_wing_samples <- c(151 + 138, 151 + 131, 151 + 129, 151 + 131)
n_left_wing_samples <- c(209 + 208, 209 + 205, 209 + 200, 209 + 199)

median_full_sample <- round(median(n_full_sample))
median_right_sample <- round(median(n_right_wing_samples))
median_left_sample <- round(median(n_left_wing_samples))

# Run power analyses looping over different effect sizes

# Generate a vector with all ATE's to loop over
ates <- seq(from = 0.1, to = 0.5, by = 0.01)

# Empty containers to store results
# For full sample
full_sample_power <- list()
# For left-wing
left_sample_power <- list()
# For right-wing
right_sample_power <- list()

# Loop over ATE's, full sample
for (i in 1:length(ates)) {
  x <- ates[i]
  design <- two_arm_designer(
    N = median_full_sample,
    assignment_prob = 0.5,
    ate = x
  )
  diagnosis <- diagnose_design(design, sims = 500, bootstrap_sims = 500)

  full_sample_power[[i]] <- diagnosis$diagnosands_df[1, 16]
}

# Loop over ATE's, right-wing sample
for (i in 1:length(ates)) {
  x <- ates[i]
  design <- two_arm_designer(
    N = median_right_sample,
    assignment_prob = 0.5,
    ate = x
  )
  diagnosis <- diagnose_design(design, sims = 500, bootstrap_sims = 500)

  right_sample_power[[i]] <- diagnosis$diagnosands_df[1, 16]
}

# Loop over ATE's, left-wing sample
for (i in 1:length(ates)) {
  x <- ates[i]
  design <- two_arm_designer(
    N = median_left_sample,
    assignment_prob = 0.5,
    ate = x
  )
  diagnosis <- diagnose_design(design, sims = 500, bootstrap_sims = 500)

  left_sample_power[[i]] <- diagnosis$diagnosands_df[1, 16]
}


full_sample_power_df <- do.call(rbind.data.frame, full_sample_power)
left_sample_power_df <- do.call(rbind.data.frame, left_sample_power)
right_sample_power_df <- do.call(rbind.data.frame, right_sample_power)
ates_2 <- as.list(ates)
ates_df <- do.call(rbind.data.frame, ates_2)

# Fix column name
colnames(full_sample_power_df) <- c("Power")
colnames(left_sample_power_df) <- c("Power")
colnames(right_sample_power_df) <- c("Power")
colnames(ates_df) <- c("ATE")

# Add a variable that identifies sample
full_sample_power_df$Sample <- "Full sample"
left_sample_power_df$Sample <- "Left-wing"
right_sample_power_df$Sample <- "Right-wing"

# Merge all of this
power_df <- rbind(full_sample_power_df, left_sample_power_df, right_sample_power_df)
power_df_final <- cbind(power_df, ates_df)

max(power_df_final$ATE)

# Make the plot
fig_d1 <- ggplot(power_df_final, aes(x = ATE, y = Power)) +
  geom_line(size = 1, color = "blue") +
  xlab("Average treatment effect (standardized)") +
  ylab("Power") +
  geom_hline(yintercept = 0.8) +
  theme_bw() +
  scale_y_continuous(breaks = c(0, 0.2, 0.4, 0.6, 0.8, 1)) +
  scale_x_continuous(breaks = c(0, 0.1, 0.2, 0.3, 0.4, 0.5)) +
  theme(plot.title = element_text(face = "bold", size = 13)) +
  facet_wrap(. ~ Sample)
# fig_d1

ggsave("plots/fig_d1.png", plot = fig_d1, width = 17.5, height = 12, units = "cm")
